'''
# CT 和 PET 类型统计
# CT 和 PET 对应关系验证
# PatientID格式统一
# CT 和 PET 类型对应
'''

import os
import pydicom
import pandas as pd
import re

# 本文件完成CT和PET的对应整理

ERROR_sum = []   # 查出数据中CT数量和PET数量不一致的病人
num = 1


# 验证数据中CT和PET对应关系通过InstanceNumber, 利用imageposition[2]或slice location来验证位置是否正确
def CTcorrespond_PET_Number():
    global num
    ERROR_Number = []  # 查出数据中【‘21703’，‘8274’】的CT和PET的InstanceNumber不一致的病人
    ERROR_Location = [] # 查出数据中CT和对应PET的Location不一致的病人， 然后人工对齐, 位置大致对应（设1像素的容错）
    origin_path = 'D:\lung_cancer\data\origin_data'
    patientList = os.listdir(origin_path)
    for patient in patientList:
        print('%d--->%s'%(num, patient))
        num = num+1

        CT_PATH = os.path.join(origin_path, patient, 'CT')
        CT_NAME = os.listdir(CT_PATH)
        CT_slice = pydicom.read_file(os.path.join(CT_PATH, CT_NAME[0]))
        PET_PATH = os.path.join(origin_path, patient, 'PET')
        PET_NAME = os.listdir(PET_PATH)
        PET_slice = pydicom.read_file(os.path.join(PET_PATH, PET_NAME[0]))
        if(CT_slice.InstanceNumber != PET_slice.InstanceNumber):
            ERROR_Number.append(patient)
        if int(CT_slice.SliceLocation) != int(PET_slice.SliceLocation):
            ERROR_Location.append(patient)
    print('ERROR Number is: ', ERROR_Number)
    print('ERROR Location is: ', ERROR_Location)



# 统计所有病人的CT类型和PET类型
def statisticType():
    global num

    patient_id = list(pd.read_csv('D:/lung_cancer/data/data.csv')['patientID'])

    origin_path = 'H:'
    # year_list = ['2008', '2009', '2010', '2011', '2012', '2013', '2014', '2015', '2016', '2017']
    year_list = ['2008']
    for year in year_list:
        days = os.listdir(os.path.join(origin_path, year))
        for day in days:
            patients = os.listdir(os.path.join(origin_path, year, day))
            for patient in patients:
                if int(patient) in patient_id:
                    CT_types = {}
                    PET_types = {}
                    print('%d--->%s' % (num, patient))
                    num = num + 1
                    patient_path = os.path.join(origin_path, year, day, patient)
                    for s in os.listdir(patient_path):
                        slice = pydicom.read_file(os.path.join(patient_path, s))
                        sd = slice.SeriesDescription
                        sm = slice.Modality
                        if sm == 'CT':
                            if sd in CT_types.keys():
                                CT_types[sd] = CT_types[sd]+1
                            else:
                                CT_types[sd] = 1
                        elif sm == 'PT':
                            if sd in PET_types:
                                PET_types[sd] = PET_types[sd]+1
                            else:
                                PET_types[sd] = 1

                    csv_name = os.path.join('D:/lung_cancer/data/seriesname', patient + '_CT.csv')
                    df = pd.DataFrame([list(CT_types.values())], columns=list(CT_types.keys()))
                    df.to_csv(csv_name, index=False)

                    csv_name2 = os.path.join('D:/lung_cancer/data/seriesname', patient + '_PET.csv')
                    df2 = pd.DataFrame([list(PET_types.values())], columns=list(PET_types.keys()))
                    df2.to_csv(csv_name2, index=False)


# 统计CT有几种
def count_ct():
    src_path = 'D:/lung_cancer/data/seriesname/'
    names = os.listdir(src_path)
    ct_list = []
    for name in names:
        id = name.split('.')[0].split('_')[1]
        if id == 'CT':
            data = pd.read_csv(os.path.join(src_path, name))
            keys = list(data.keys())
            print(name.split('.')[0], keys)
            for key in keys:
                if (key not in ct_list) and (not re.match('ScreenCap', key)):
                    ct_list.append(key)
    print('type: ', ct_list)
    print('ct count: ', len(ct_list))

# 统计pet有几种
def count_pet():
    src_path = 'D:/lung_cancer/data/seriesname/'
    names = os.listdir(src_path)
    pet_list = []
    for name in names:
        id = name.split('.')[0].split('_')[1]
        if id == 'PET':
            data = pd.read_csv(os.path.join(src_path, name))
            keys = list(data.keys())
            print(name.split('.')[0], keys)
            for key in keys:
                if (key not in pet_list) and (not re.match('ScreenCap', key)):
                    pet_list.append(key)
    print('type: ', pet_list)
    print('pet count: ', len(pet_list))


# 部分病人id有：08988， 28674-2这种不规则形式，统一修改为能转化为int型数据的样式，这里统计不规则样式
def count_abnormal_format():
    abnormal_list = []
    origin_path = 'H:'
    year_list = ['2008', '2009', '2010', '2011', '2012', '2013', '2014', '2015', '2016', '2017']
    for year in year_list:
        days = os.listdir(os.path.join(origin_path, year))
        for day in days:
            patients = os.listdir(os.path.join(origin_path, year, day))
            for patient in patients:
                if not patient.isdigit() or patient[0]=='0':
                    abnormal_list.append(year+'----'+day+'----'+patient)
    print(len(abnormal_list))
    print(abnormal_list)


def count_923_1141():
    global num

    patient_id = list(pd.read_csv('D:/lung_cancer/data/data.csv')['patientID'])

    name_list = []

    origin_path = 'H:'
    year_list = ['2008', '2009', '2010', '2011', '2012', '2013', '2014', '2015', '2016', '2017']
    for year in year_list:
        days = os.listdir(os.path.join(origin_path, year))
        for day in days:
            patients = os.listdir(os.path.join(origin_path, year, day))
            for patient in patients:
                if int(patient) in patient_id:
                    if int(patient) not in name_list:
                        name_list.append(int(patient))
                    else:
                        print(int(patient))
    print(name_list)
    print(len(name_list))

# 统计CT或PET类别多于一种的病人，进行手工对应
def count_more_one_patient():
    global num
    patient_id = list(pd.read_csv('D:/lung_cancer/data/data.csv')['patientID'])
    src_path = 'D:/lung_cancer/data/seriesname/'
    for id in patient_id:
        pet_name = os.path.join(src_path+str(id)+'_PET.csv')
        data = pd.read_csv(pet_name)
        pet_key_list = []
        for key in list(data.keys()):
            if not re.match('ScreenCap', key):
                pet_key_list.append(key)

        ct_name = os.path.join(src_path + str(id) + '_CT.csv')
        data2 = pd.read_csv(ct_name)
        ct_key_list = []
        for key2 in list(data2.keys()):
            if not re.match('ScreenCap', key2):
                ct_key_list.append(key2)
        if len(pet_key_list) != 1 or len(ct_key_list) != 1:
            print(num, '--->', id, 'ct type: ', ct_key_list, 'pet type: ', pet_key_list)
            num = num+1

if __name__ == '__main__':
    # statisticType()
    # count_923_1141()
    # count_ct()
    # count_pet()
    # count_more_one_patient()
    # count_abnormal_format()
    CTcorrespond_PET_Number()
